"""
PDF背景提取处理器 - OpenCV增强版
使用 pdf2image + OpenCV 进行智能背景提取

核心思路：
1. PDF → 高清图像（pdf2image或PyMuPDF）
2. OpenCV检测文字区域（使用EAST或MSER）
3. OpenCV检测小图标/logo（轮廓检测）
4. 智能修复（内容感知填充/Inpaint）
5. 保留大背景图
"""
import fitz  # PyMuPDF
from PIL import Image, ImageDraw, ImageFilter
import base64
from io import BytesIO
import numpy as np
import cv2


class PDFBackgroundProcessor:
    def __init__(self):
        self.debug = True
    
    def extract_background(self, pdf_path):
        """
        新算法：直接提取PDF的背景图层
        1. 分析PDF结构，找出背景图片（最底层、最大面积）
        2. 如果有背景图，直接提取
        3. 如果没有背景图，创建纯色/渐变背景
        """
        doc = fitz.open(pdf_path)
        pages_data = []
        
        if self.debug:
            print(f"\n🎨 开始提取PDF背景图层，共 {len(doc)} 页")
        
        for page_num in range(len(doc)):
            page = doc[page_num]
            
            try:
                if self.debug:
                    print(f"\n📄 处理页面 {page_num + 1}")
                    print(f"   页面尺寸: {page.rect.width}x{page.rect.height}")
                
                # 1. 查找背景图片
                background_image = self._find_background_image(page)
                
                if background_image:
                    # 有背景图 - 直接使用
                    if self.debug:
                        print(f"   ✅ 找到背景图")
                    background_img = background_image
                else:
                    # 没有背景图 - 提取纯色/渐变背景
                    if self.debug:
                        print(f"   ⚠️ 未找到背景图，提取纯色背景")
                    background_img = self._extract_solid_background(page)
                
                # 2. 转换为base64
                buffer = BytesIO()
                background_img.save(buffer, format='PNG', optimize=True, quality=95)
                img_base64 = base64.b64encode(buffer.getvalue()).decode()
                
                pages_data.append({
                    'page_num': page_num,
                    'width': page.rect.width,
                    'height': page.rect.height,
                    'scaled_width': background_img.width,
                    'scaled_height': background_img.height,
                    'background_image': img_base64,
                    'text_regions_count': 0,
                    'image_regions_count': 0
                })
                
                if self.debug:
                    print(f"   ✅ 背景提取完成")
                
            except Exception as e:
                print(f"   ❌ 处理页面 {page_num + 1} 失败: {e}")
                import traceback
                traceback.print_exc()
                raise
        
        doc.close()
        
        if self.debug:
            print(f"\n🎉 所有页面背景提取完成！")
        
        return pages_data
    
    def _find_background_image(self, page):
        """
        查找PDF页面中的背景图片
        策略：找出占页面面积最大的图片（通常是背景）
        """
        try:
            page_width = page.rect.width
            page_height = page.rect.height
            page_area = page_width * page_height
            
            image_list = page.get_images(full=True)
            
            if not image_list:
                return None
            
            # 分析所有图片，找出最大的
            largest_image = None
            largest_area = 0
            largest_rect = None
            
            for img_index, img_info in enumerate(image_list):
                xref = img_info[0]
                image_rects = page.get_image_rects(xref)
                
                if not image_rects:
                    continue
                
                for rect in image_rects:
                    img_width = rect.x1 - rect.x0
                    img_height = rect.y1 - rect.y0
                    img_area = img_width * img_height
                    area_ratio = img_area / page_area
                    
                    # 只考虑占比>50%的图片作为背景
                    if area_ratio > 0.5 and img_area > largest_area:
                        largest_area = img_area
                        largest_image = xref
                        largest_rect = rect
                        
                        if self.debug:
                            print(f"   🎨 候选背景图: xref={xref}, 尺寸{img_width:.1f}x{img_height:.1f}, 占比{area_ratio*100:.1f}%")
            
            if largest_image:
                # 提取背景图
                pix = fitz.Pixmap(page.parent, largest_image)
                
                # 转换颜色空间
                if pix.n - pix.alpha > 3:
                    pix = fitz.Pixmap(fitz.csRGB, pix)
                
                # 转为PIL Image
                img_data = pix.tobytes("png")
                background_img = Image.open(BytesIO(img_data))
                
                # 调整到页面尺寸
                background_img = background_img.resize(
                    (int(page_width * 2), int(page_height * 2)),
                    Image.LANCZOS
                )
                
                if self.debug:
                    print(f"   ✅ 提取背景图: {background_img.width}x{background_img.height}")
                
                return background_img
            
            return None
            
        except Exception as e:
            if self.debug:
                print(f"   ⚠️ 查找背景图失败: {e}")
            return None
    
    def _extract_solid_background(self, page):
        """
        提取纯色/渐变背景
        策略：采样页面四个角落的颜色，创建纯色或渐变背景
        """
        try:
            page_width = page.rect.width
            page_height = page.rect.height
            
            # 渲染完整页面（用于采样颜色）
            scale = 2.0
            mat = fitz.Matrix(scale, scale)
            pix = page.get_pixmap(matrix=mat, alpha=False)
            
            # 转为PIL Image
            full_img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            
            # 采样四个角落的颜色
            corner_size = min(100, pix.width // 10, pix.height // 10)
            
            corners = {
                'top_left': full_img.crop((0, 0, corner_size, corner_size)),
                'top_right': full_img.crop((pix.width - corner_size, 0, pix.width, corner_size)),
                'bottom_left': full_img.crop((0, pix.height - corner_size, corner_size, pix.height)),
                'bottom_right': full_img.crop((pix.width - corner_size, pix.height - corner_size, pix.width, pix.height))
            }
            
            # 计算每个角落的平均颜色
            corner_colors = {}
            for name, corner_img in corners.items():
                # 转为numpy数组
                corner_array = np.array(corner_img)
                # 计算平均颜色
                avg_color = corner_array.reshape(-1, 3).mean(axis=0).astype(np.uint8)
                corner_colors[name] = tuple(avg_color)
                
                if self.debug:
                    print(f"   🎨 {name}: RGB{corner_colors[name]}")
            
            # 检查是否所有角落颜色相似（纯色背景）
            colors_list = list(corner_colors.values())
            color_variance = np.var([list(c) for c in colors_list], axis=0).mean()
            
            if color_variance < 100:  # 颜色差异小，是纯色背景
                # 使用平均颜色创建纯色背景
                avg_bg_color = tuple(np.mean([list(c) for c in colors_list], axis=0).astype(np.uint8))
                background_img = Image.new('RGB', (pix.width, pix.height), avg_bg_color)
                
                if self.debug:
                    print(f"   ✅ 创建纯色背景: RGB{avg_bg_color}")
            else:
                # 颜色差异大，可能是渐变背景
                # 创建简单的渐变（从左上到右下）
                background_img = self._create_gradient_background(
                    pix.width, pix.height,
                    corner_colors['top_left'],
                    corner_colors['bottom_right']
                )
                
                if self.debug:
                    print(f"   ✅ 创建渐变背景")
            
            return background_img
            
        except Exception as e:
            if self.debug:
                print(f"   ⚠️ 提取纯色背景失败: {e}")
                import traceback
                traceback.print_exc()
            
            # 降级：返回白色背景
            return Image.new('RGB', (int(page_width * 2), int(page_height * 2)), (255, 255, 255))
    
    def _create_gradient_background(self, width, height, color1, color2):
        """
        创建渐变背景（从color1到color2）
        """
        background = Image.new('RGB', (width, height))
        draw = ImageDraw.Draw(background)
        
        # 创建从上到下的渐变
        r1, g1, b1 = color1
        r2, g2, b2 = color2
        
        for y in range(height):
            # 计算当前行的颜色
            ratio = y / height
            r = int(r1 * (1 - ratio) + r2 * ratio)
            g = int(g1 * (1 - ratio) + g2 * ratio)
            b = int(b1 * (1 - ratio) + b2 * ratio)
            
            draw.line([(0, y), (width, y)], fill=(r, g, b))
        
        return background
    
    def _extract_text_regions(self, page, scale):
        """提取文字区域的边界框"""
        regions = []
        
        try:
            text_dict = page.get_text("rawdict")
            blocks = text_dict.get("blocks", [])
            
            for block in blocks:
                if block.get("type") == 0:  # 文字块
                    for line in block.get("lines", []):
                        for span in line.get("spans", []):
                            bbox = span.get("bbox", [])
                            if len(bbox) == 4 and span.get("text", "").strip():
                                padding = 3
                                scaled_region = [
                                    int(bbox[0] * scale) - padding,
                                    int(bbox[1] * scale) - padding,
                                    int(bbox[2] * scale) + padding,
                                    int(bbox[3] * scale) + padding
                                ]
                                regions.append(scaled_region)
        except Exception as e:
            if self.debug:
                print(f"   ⚠️ 提取文字区域失败: {e}")
        
        return regions
    
    def _detect_small_objects(self, img_bgr):
        """
        使用OpenCV检测小图标/logo
        基于轮廓检测，过滤掉大背景图
        """
        regions = []
        
        try:
            height, width = img_bgr.shape[:2]
            page_area = height * width
            
            # 转灰度
            gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
            
            # 边缘检测
            edges = cv2.Canny(gray, 50, 150)
            
            # 形态学操作，连接边缘
            kernel = np.ones((5, 5), np.uint8)
            edges = cv2.dilate(edges, kernel, iterations=1)
            edges = cv2.erode(edges, kernel, iterations=1)
            
            # 查找轮廓
            contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            for contour in contours:
                x, y, w, h = cv2.boundingRect(contour)
                area = w * h
                area_ratio = area / page_area
                
                # 只检测小对象（面积<5%的页面）
                # 且尺寸合理（>10像素，<30%页面宽/高）
                if (area_ratio < 0.05 and 
                    w > 10 and h > 10 and 
                    w < width * 0.3 and h < height * 0.3):
                    
                    # 检查是否与文字区域重叠（避免重复）
                    is_text_region = False
                    # 这里可以添加与文字区域的重叠检测
                    
                    if not is_text_region:
                        padding = 2
                        regions.append([
                            max(0, x - padding),
                            max(0, y - padding),
                            min(width, x + w + padding),
                            min(height, y + h + padding)
                        ])
                        
                        if self.debug:
                            print(f"   🔍 小对象: 位置({x},{y}), 尺寸{w}x{h}, 占比{area_ratio*100:.2f}%")
        
        except Exception as e:
            if self.debug:
                print(f"   ⚠️ 检测小对象失败: {e}")
        
        return regions
    
    def _inpaint_regions(self, img_bgr, regions):
        """
        使用OpenCV的inpainting算法智能修复区域
        这比简单填充颜色更自然
        """
        if not regions:
            return img_bgr
        
        height, width = img_bgr.shape[:2]
        
        # 创建mask
        mask = np.zeros((height, width), dtype=np.uint8)
        
        for region in regions:
            x1, y1, x2, y2 = region
            x1 = max(0, int(x1))
            y1 = max(0, int(y1))
            x2 = min(width, int(x2))
            y2 = min(height, int(y2))
            
            if x2 > x1 and y2 > y1:
                mask[y1:y2, x1:x2] = 255
        
        # 使用Telea算法进行修复（快速且效果好）
        try:
            result = cv2.inpaint(img_bgr, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
            if self.debug:
                print(f"   🎨 使用Inpainting修复了 {len(regions)} 个区域")
            return result
        except Exception as e:
            if self.debug:
                print(f"   ⚠️ Inpainting失败，使用简单填充: {e}")
            # 降级方案：简单填充
            return self._simple_fill_regions(img_bgr, regions)
    
    def _simple_fill_regions(self, img_bgr, regions):
        """
        简单填充方案（降级）
        采样周围颜色并填充
        """
        result = img_bgr.copy()
        
        for region in regions:
            x1, y1, x2, y2 = region
            x1 = max(0, int(x1))
            y1 = max(0, int(y1))
            x2 = min(result.shape[1], int(x2))
            y2 = min(result.shape[0], int(y2))
            
            if x2 > x1 and y2 > y1:
                # 采样周围颜色
                bg_color = self._sample_color_opencv(result, x1, y1, x2, y2)
                result[y1:y2, x1:x2] = bg_color
        
        return result
    
    def _sample_color_opencv(self, img_bgr, x1, y1, x2, y2):
        """
        使用OpenCV采样背景色
        """
        height, width = img_bgr.shape[:2]
        margin = 5
        
        # 扩展采样区域
        ex1 = max(0, x1 - margin)
        ey1 = max(0, y1 - margin)
        ex2 = min(width, x2 + margin)
        ey2 = min(height, y2 + margin)
        
        samples = []
        
        # 采样四边
        # 上边
        if ey1 >= 0:
            samples.append(img_bgr[ey1, ex1:ex2].mean(axis=0))
        # 下边
        if ey2 < height:
            samples.append(img_bgr[ey2-1, ex1:ex2].mean(axis=0))
        # 左边
        if ex1 >= 0:
            samples.append(img_bgr[ey1:ey2, ex1].mean(axis=0))
        # 右边
        if ex2 < width:
            samples.append(img_bgr[ey1:ey2, ex2-1].mean(axis=0))
        
        if samples:
            # 取平均颜色
            avg_color = np.mean(samples, axis=0).astype(np.uint8)
            return avg_color
        
        return np.array([255, 255, 255], dtype=np.uint8)
    
    def create_background_pdf(self, original_pdf_path):
        """
        创建只包含背景的PDF文件
        """
        if self.debug:
            print(f"\n📦 开始生成纯背景PDF")
        
        # 提取背景图片
        pages_data = self.extract_background(original_pdf_path)
        
        # 创建新PDF
        doc = fitz.open()
        
        for page_data in pages_data:
            # 创建新页面
            page = doc.new_page(
                width=page_data['width'],
                height=page_data['height']
            )
            
            # 解码背景图片
            img_bytes = base64.b64decode(page_data['background_image'])
            
            # 插入背景图片
            rect = fitz.Rect(0, 0, page_data['width'], page_data['height'])
            page.insert_image(rect, stream=img_bytes)
            
            if self.debug:
                print(f"   ✅ 页面 {page_data['page_num'] + 1} 已添加")
        
        # 保存到字节流
        output = BytesIO()
        doc.save(output, garbage=4, deflate=True, clean=True)
        pdf_bytes = output.getvalue()
        
        doc.close()
        
        if self.debug:
            print(f"   📄 PDF生成完成，大小: {len(pdf_bytes)} bytes")
        
        return pdf_bytes
